package utils

import (
	"crypto/ecdsa"
	"crypto/ed25519"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"encoding/pem"
	"errors"

	"golang.org/x/crypto/ssh"
)

func GeneratePrivateKey(algo string, comment string) ([]byte, error) {
	var pemBlock *pem.Block
	if algo == ssh.KeyAlgoRSA {
		signer, err := rsa.GenerateKey(rand.Reader, 4096)
		if err != nil {
			return nil, err
		}
		pemBlock, err = ssh.MarshalPrivateKey(signer, comment)
		if err != nil {
			return nil, err
		}
	} else if algo == ssh.KeyAlgoED25519 {
		_, signer, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return nil, err
		}
		pemBlock, err = ssh.MarshalPrivateKey(signer, comment)
		if err != nil {
			return nil, err
		}
	} else if algo == ssh.KeyAlgoECDSA256 {
		signer, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		if err != nil {
			return nil, err
		}
		pemBlock, err = ssh.MarshalPrivateKey(signer, comment)
		if err != nil {
			return nil, err
		}
	} else {
		return nil, errors.New("unsupported algorithm")
	}

	return pem.EncodeToMemory(pemBlock), nil

}
